#include "RemoteThreadHelper.h"
#include <psapi.h>
#include <tchar.h>

HMODULE CRemoteProcess::Load(LPCTSTR pszName)
{
	return (HMODULE)m_rfLoadLibrary.Execute((LPVOID)pszName);	//Pay attention in Win64
}

BOOL CRemoteProcess::Free(HMODULE hmod)
{
	return m_rfFreeLibrary.Execute(hmod);
}

HMODULE CRemoteModuleBase::FindProcessModule(HANDLE hProcess, LPCTSTR pszBaseName)
{
	HMODULE* phmods = NULL;
	DWORD cbAlloc = 0;
	DWORD cbReturned = 0;

	do
	{
		if ((cbAlloc += sizeof(HMODULE) * 1024) < cbReturned)
			break;

		HMODULE* phmod = (HMODULE*)realloc(phmods, cbAlloc);
		if (!phmod)
			break;

		phmods = phmod;
		EnumProcessModules(hProcess, phmods, cbAlloc, &cbReturned);
	}
	while (cbReturned == cbAlloc);

	if (!phmods) return NULL;

	HMODULE hmod = NULL;
	for (DWORD i = 0; i < cbReturned / sizeof(HMODULE); i++)
	{
		TCHAR szName[MAX_PATH + 1];
		*szName = 0;
		GetModuleBaseName(hProcess, phmods[i], szName, sizeof(szName) / sizeof(szName[0]));
		if (!_tcsicmp(szName, pszBaseName))
		{
			hmod = phmods[i];
			break;
		}
	}

	free(phmods);
	return hmod;
}

BOOL CRemoteFunction::Create(CRemoteModuleBase* pModule, LPCSTR pszName)
{
	if(pModule)
	{
		m_prm = pModule;
		m_pfnLocal = GetProcAddress(m_prm->GetLocalHandle(), pszName);
		return m_pfnLocal != NULL;
	}

	return FALSE;
}

DWORD CRemoteFunction::Execute(LPVOID pParameter)
{
	if((NULL == m_prm)||(NULL == (HMODULE)m_prm)) return 0;

	DWORD dwThreadId = 0;
	HANDLE hThread = CreateRemoteThread(*(m_prm->GetRemoteProcess()), NULL, 0, 
		(LPTHREAD_START_ROUTINE)((PBYTE)m_pfnLocal - (PBYTE)m_prm->GetLocalHandle() + (PBYTE)(HMODULE)*m_prm), 
		pParameter, 0, &dwThreadId);

	if (!hThread) return 0;

	WaitForSingleObject(hThread, INFINITE);

	DWORD dwExitCode;
	GetExitCodeThread(hThread, &dwExitCode);

	CloseHandle(hThread);
	return dwExitCode;
}

BOOL CRemoteModuleNL::Create(CRemoteProcess* prp, LPCTSTR pszBaseName)
{
	if(NULL == prp) return FALSE;

	HMODULE hmodL = GetModuleHandle(pszBaseName);
	if (!hmodL)
		return FALSE;

	HMODULE hmod = FindProcessModule((*prp), pszBaseName);
	if (!hmod)
		return FALSE;

	m_prp = prp;
	m_h = hmod;
	m_hLocal = hmodL;
	return TRUE;
}

BOOL CRemoteProcess::FindProcess(LPCTSTR pszBaseName, DWORD dwDesiredAccess, 
								   HANDLE* pHandle, DWORD& dwHandleCount)
{
	if(pHandle == NULL) return FALSE;

	DWORD* pdwPIDs = NULL;
	DWORD cbAlloc = 0;
	DWORD cbReturned = 0;
	
	dwHandleCount = 0;

	do
	{
		if ((cbAlloc += sizeof(DWORD) * 1024) < cbReturned)
			break;

		DWORD* pdw = (DWORD*)realloc(pdwPIDs, cbAlloc);
		if (!pdw)
			break;
		pdwPIDs = pdw;
		EnumProcesses(pdwPIDs, cbAlloc, &cbReturned);
	}
	while (cbReturned == cbAlloc);

	if (!pdwPIDs) return NULL;

	DWORD dwCurSessionId = 0;
	ProcessIdToSessionId(GetCurrentProcessId(), &dwCurSessionId);

	DWORD dwSessionId = 0;
	HANDLE hProcess = NULL;
	for (DWORD i = 0; i < cbReturned / sizeof(DWORD); i++)
	{
		ProcessIdToSessionId(pdwPIDs[i], &dwSessionId);
		if(dwCurSessionId != dwSessionId) continue;

		hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | dwDesiredAccess, FALSE, pdwPIDs[i]);
		if (!hProcess)
			continue;

		TCHAR szName[MAX_PATH + 1];
		*szName = 0;
		GetProcessImageFileName(hProcess, szName, sizeof(szName) / sizeof(szName[0]));

		TCHAR* p = _tcsrchr(szName, '\\');
		p = p ? p + 1: szName;
		if (!_tcsicmp(p, pszBaseName))
		{
			pHandle[dwHandleCount] = hProcess;
			dwHandleCount++;
		}
		else
		{
			CloseHandle(hProcess);
			hProcess = NULL;
		}
	}

	free(pdwPIDs);

	return (dwHandleCount>0);
}

#ifdef UNICODE
#define FN_LoadLibrary "LoadLibraryW"
#else
#define FN_LoadLibrary "LoadLibraryA"
#endif

HANDLE CRemoteProcess::GetTargetProcess(HANDLE* ph, DWORD dwHandleCount)
{
	if(ph == NULL) return NULL;
	
	return *ph;
}

BOOL CRemoteProcess::Create(LPCTSTR pszBaseName, DWORD dwDesiredAccess)
{
	HANDLE h[16];
	DWORD dwHandleCount = 0;

	BOOL bFind = FindProcess(pszBaseName, dwDesiredAccess, h, dwHandleCount);
	if (bFind)
	{
		m_h = GetTargetProcess(h, dwHandleCount);
		if (!m_rmKernel.Create(this, _T("kernel32.dll"))) goto _end;
		if (!m_rfLoadLibrary.Create(&m_rmKernel, FN_LoadLibrary)) goto _end;
		if (!m_rfFreeLibrary.Create(&m_rmKernel, "FreeLibrary")) goto _end;
		return TRUE;
	}
	
_end:
	if(m_h)
	{
		CloseHandle(m_h);
		m_h = NULL;
	}

	return FALSE;
}

HANDLE CRemoteProcessIE::GetTargetProcess(HANDLE* ph, DWORD dwHandleCount)
{
	HANDLE h = CRemoteProcess::GetTargetProcess(ph, dwHandleCount);
	if(h&&dwHandleCount > 1)
	{
		CRemoteModuleBase rm(this);

		for(DWORD i = 0; i<dwHandleCount; i++)
		{
			if(rm.FindProcessModule(ph[i], _T("D3D9.dll")))
			{
				h = ph[i];
			}
			else
			{
				CloseHandle(ph[i]);
				ph[i] = 0;
			}
		}
	}

	return h;
}

BOOL CRemoteModule::Create(LPCTSTR pszName, HMODULE hLocal, HMODULE hRemote)
{
	HMODULE hmodL = NULL, hmod = NULL;

	if(hLocal)
		hmodL = hLocal;
	else
	{
		if(pszName) hmodL = LoadLibrary(pszName);
	}

	if (!hmodL) return FALSE;

	if(hRemote)
		hmod = hRemote;
	else
	{
		if(NULL == pszName) return FALSE;

		int nSize = sizeof(*pszName) * (_tcslen(pszName) + 1);
		LPTSTR pszNameR = (LPTSTR)VirtualAllocEx(*m_prp, NULL, nSize, MEM_COMMIT, PAGE_READWRITE);
		if (!pszNameR)
		{
			FreeLibrary(hmodL);
			return FALSE;
		}

		if(WriteProcessMemory(*m_prp, pszNameR, pszName, nSize, NULL))
			hmod = m_prp->Load(pszNameR);

		VirtualFreeEx(*m_prp, pszNameR, 0, MEM_RELEASE);
		if (!hmod)
		{
			FreeLibrary(hmodL);
			return FALSE;
		}
	}

	m_h = hmod;
	m_hLocal = hmodL;

	return TRUE;
}

void CRemoteModule::Destroy()
{
	if (m_h)
	{
		m_prp->Free(m_h);
		m_h = NULL;
	}

	if(m_hLocal)
	{
		FreeLibrary(m_hLocal);
		m_hLocal = NULL;
	}
}